#include "StdAfx.h"
#include "OverlappedCompletionRoutineModel.h"
#include "IAsyncSendQueue.h"

INetworkLib* OverlappedCompletionRoutineModelLib::CreateComponent()
{
	return new COverlappedCompletionRoutineModel;
}

struct WorkerThreadDesc
{
	SOCKET				m_socket;
	COverlappedCompletionRoutineModel*	m_pNetworkLib;	
	WorkerThreadDesc(SOCKET m_socket
		, COverlappedCompletionRoutineModel& networkLib
		)
		: m_socket(m_socket)
		, m_pNetworkLib(&networkLib)					
	{

	}
};

struct WSAOVERLAPPEDEX : public WSAOVERLAPPED
{
	COverlappedCompletionRoutineModel*	m_pNetworkLib;
	SOCKET								socket;
	char								buffer[0xff];
};

COverlappedCompletionRoutineModel::COverlappedCompletionRoutineModel(void)
	: m_pAsyncSendQueue(nullptr)
{
}


COverlappedCompletionRoutineModel::~COverlappedCompletionRoutineModel(void)
{
	if (m_pAsyncSendQueue)
	{
		m_pAsyncSendQueue->Release();
		m_pAsyncSendQueue = nullptr;
	}
}

bool COverlappedCompletionRoutineModel::Initialize( ON_CONNECT_FUNCTOR onConnect , ON_ACCEPT_FUNCTOR onAccept , ON_DISCONNECT_FUNCTOR onDisconnect , ON_RECV_FUNCTOR onRecv )
{
	if (false == Helper::WSAStartup())
	{
		return false;
	}

	m_onConnectFunctor		= onConnect;
	m_onAcceptFunctor		= onAccept;
	m_onDisconnectFunctor	= onDisconnect;
	m_onRecvFunctor			= onRecv;

	m_pAsyncSendQueue = AsyncSendQueueForOverlappedModel::CreateComponent(
		[&](SOCKET socket){
			Disconnect(socket);
	}
	, [&](SOCKET socket)->WSAOVERLAPPED*{
		WSAOVERLAPPEDEX* pOverlapped = new WSAOVERLAPPEDEX;
		memset(pOverlapped, 0, sizeof(WSAOVERLAPPED));
		pOverlapped->socket = socket;
		pOverlapped->m_pNetworkLib = this;
		return pOverlapped;
	}
	, sendCompletionRoutine);

	return true;
}

bool COverlappedCompletionRoutineModel::Connect( const std::tstring& ip, WORD port )
{
	SOCKET connectSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (connectSocket == INVALID_SOCKET) 
	{
		return false;
	}
	
	// Connect to server.
	SOCKADDR_IN sockaddr;	
	INT addressSize = 0;	
	Helper::WSAStringToAddress(ip, port, sockaddr, addressSize);	

	int iResult = connect( connectSocket, (SOCKADDR*)&sockaddr, addressSize);
	if (iResult == SOCKET_ERROR) 
	{
		if (WSAGetLastError() != WSAEWOULDBLOCK)
		{
			return false;
		}		
	}	

	if (false == postRecv(connectSocket, nullptr))
	{
		Disconnect(connectSocket);
		return false;
	}

	m_onConnectFunctor(connectSocket, 0);

	return true;
}

bool COverlappedCompletionRoutineModel::Send( SOCKET socket, IStreamBuffer& stream )
{
	return m_pAsyncSendQueue->Send(socket, stream);
}

bool COverlappedCompletionRoutineModel::Send( SOCKET socket, size_t bytesSent )
{
	return m_pAsyncSendQueue->Send(socket, bytesSent);
}

bool COverlappedCompletionRoutineModel::Disconnect( SOCKET socket )
{
	shutdown(socket, SD_SEND);
	closesocket(socket);
	m_onDisconnectFunctor(socket);
	return true;
}

bool COverlappedCompletionRoutineModel::Listen( const std::tstring& ip, WORD port, int backlog )
{
	SOCKET listenSocket = Helper::WSAListen(ip, port, backlog);
	if (listenSocket == INVALID_SOCKET)
	{
		return false;
	}

	WorkerThreadDesc* pArgment = new WorkerThreadDesc(listenSocket, *this);

	UINT threadID = 0;
	HANDLE hThread = (HANDLE)_beginthreadex(0, 0, acceptWorkerThread, pArgment, 0, &threadID);
	if (hThread == INVALID_HANDLE_VALUE)
	{
		return false;
	}	

	CloseHandle(hThread);

	return true;
}

void COverlappedCompletionRoutineModel::Release()
{
	Helper::WSACleanup();
	delete this;
}

void CALLBACK COverlappedCompletionRoutineModel::sendCompletionRoutine(DWORD Error, 
	DWORD BytesTransferred, 
	LPWSAOVERLAPPED pOverlapped,
	DWORD InFlags)
{
	WSAOVERLAPPEDEX* pOverlappedEx = (WSAOVERLAPPEDEX*)pOverlapped;

	if (Error != 0 || BytesTransferred == 0)
	{
		// Either a bad error occurred on the socket
		// or the socket was closed by a peer
		pOverlappedEx->m_pNetworkLib->Disconnect(pOverlappedEx->socket);
		
		delete pOverlappedEx;
		return;
	}
	
	pOverlappedEx->m_pNetworkLib->Send(pOverlappedEx->socket, BytesTransferred);
}

void CALLBACK COverlappedCompletionRoutineModel::recvCompletionRoutine(DWORD Error, 
	DWORD BytesTransferred, 
	LPWSAOVERLAPPED pOverlapped,
	DWORD InFlags)
{
	WSAOVERLAPPEDEX* pOverlappedEx = (WSAOVERLAPPEDEX*)pOverlapped;
	
	if (Error != 0 || BytesTransferred == 0)
	{
		// Either a bad error occurred on the socket
		// or the socket was closed by a peer
		pOverlappedEx->m_pNetworkLib->Disconnect(pOverlappedEx->socket);
		
		delete pOverlappedEx;
		return;
	}

	IStreamBuffer* pStreamBuffer = StreamBuffer::CreateComponent();
	pStreamBuffer->Write(pOverlappedEx->buffer, BytesTransferred);	
	pOverlappedEx->m_pNetworkLib->m_onRecvFunctor(pOverlappedEx->socket, pStreamBuffer);

	pOverlappedEx->m_pNetworkLib->postRecv(pOverlappedEx->socket, pOverlappedEx);
}

UINT __stdcall COverlappedCompletionRoutineModel::acceptWorkerThread( void* pArguments )
{
	WorkerThreadDesc* pWorkerThreadDesc = (WorkerThreadDesc*)pArguments;

	while(TRUE)
	{
		SOCKET clientSocket = accept(pWorkerThreadDesc->m_socket,NULL, NULL);

		if (clientSocket == INVALID_SOCKET)
		{
			continue;
		}
		
		if (false == pWorkerThreadDesc->m_pNetworkLib->postRecv(clientSocket, nullptr))
		{
			continue;
		}

		pWorkerThreadDesc->m_pNetworkLib->m_onAcceptFunctor(clientSocket);
	}

	delete pWorkerThreadDesc;

	return 0;
}

bool COverlappedCompletionRoutineModel::postRecv( SOCKET clientSocket, WSAOVERLAPPEDEX* pOverlapped )
{
	if (pOverlapped == nullptr)
	{	
		pOverlapped = new WSAOVERLAPPEDEX;
	}

	memset(pOverlapped, 0, sizeof(WSAOVERLAPPED));
	pOverlapped->socket = clientSocket;
	pOverlapped->m_pNetworkLib = this;

	WSABUF DataBuf;
	DataBuf.len = sizeof(pOverlapped->buffer);
	DataBuf.buf = pOverlapped->buffer;

	DWORD RecvBytes = 0;
	DWORD Flags = 0;

	if (WSARecv(clientSocket, &DataBuf, 1, &RecvBytes, &Flags, pOverlapped, recvCompletionRoutine) == SOCKET_ERROR)
	{
		if (WSAGetLastError() != WSA_IO_PENDING)
		{
			// Error occurred
			Disconnect(clientSocket);
			delete pOverlapped;

			return false;
		}
	}

	return true;
}
